#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File     : loader.py
@Project  : pipecoco
@Date     : 2021/8/23
@Author   : Zhang Jinyang
@Contact  : zhang-jy@sjtu.edu.cn
'''
 
from mindspore import load_checkpoint, load_param_into_net
import os
from net.vgg16.vgg import vgg16
from net.vgg16.config import small_vgg_cifar_cfg, large_vgg_cifar_cfg
from pipecoco.utils import dfs_search_layer
import logging

logger = logging.getLogger("net")

def load_vgg_small(model_path = None):

    cfg = small_vgg_cifar_cfg
    network = vgg16(cfg.num_classes, cfg)

    if model_path!=None and os.path.exists(model_path):
        param_dict = load_checkpoint(model_path)
        load_param_into_net(network, param_dict)
        logger.info("{} load checkpoint from [{}].".format(network.cls_name,model_path))
    network.set_train(False)

    layers_list = []
    dfs_search_layer(network, layers_list)
    return layers_list

def load_vgg_large(model_path = None):

    cfg = large_vgg_cifar_cfg
    network = vgg16(cfg.num_classes, cfg)

    if model_path!=None and os.path.exists(model_path):
        param_dict = load_checkpoint(model_path)
        load_param_into_net(network, param_dict)
        logger.info("{} load checkpoint from [{}].".format(network.cls_name,model_path))
    network.set_train(False)

    layers_list = []
    dfs_search_layer(network, layers_list)
    return layers_list